package com.service.data.flume.tools.rocketmq.sinks

import com.service.data.commons.utils.Preconditions
import org.apache.flume.Sink.Status
import org.apache.flume.conf.Configurable
import org.apache.flume.sink.AbstractSink
import org.apache.flume.{Context, Event}
import org.apache.rocketmq.client.exception.MQClientException
import org.apache.rocketmq.client.producer.{DefaultMQProducer, SendCallback, SendResult}
import org.apache.rocketmq.common.message.Message

/**
  * @author 伍鲜
  *
  *         Flume Sink to RocketMQ
  */
class RocketMQSink extends AbstractSink with Configurable {
  // 必填项
  private var namesrvAddr: String = _
  private var sendMessageMode: String = _

  // 建议填
  private var producerGroup: String = _
  private var topic: String = _

  // 可选项
  private var tags: String = ""
  private var keys: String = ""
  private var flag: Int = 0
  private var waitStoreMsgOK: Boolean = true
  private var retryTimesWhenSendAsyncFailed: Int = 2

  private var producer: DefaultMQProducer = _

  /**
    * 配置，环境初始化
    *
    * @param context
    */
  override def configure(context: Context): Unit = {
    // 必填项
    namesrvAddr = Preconditions.checkNotNull(context.getString("namesrvAddr"), "Rocket MQ namesrvAddr must be specified. For example: agent.sinks.sink.namesrvAddr=127.0.0.1:9876")
    sendMessageMode = Preconditions.checkNotNull(context.getString("sendMessageMode"), "Rocket MQ sendMessageMode must be specified. For example: agent.sinks.sink.sendMessageMode=sync").toLowerCase

    // 建议项
    producerGroup = Preconditions.getOrElse(context.getString("producerGroup"), "DEFAULT_PRODUCER_GROUP")
    topic = Preconditions.getOrElse(context.getString("topic"), "DEFAULT_TOPIC")

    // 可选项
    tags = context.getString("tags", "")
    keys = context.getString("keys", "")
    flag = context.getInteger("flag", 0)
    waitStoreMsgOK = context.getBoolean("waitStoreMsgOK", true)
    retryTimesWhenSendAsyncFailed = context.getInteger("retryTimesWhenSendAsyncFailed", 2)

    producer = new DefaultMQProducer(producerGroup)
    producer.setNamesrvAddr(namesrvAddr)

    producer.start()
    if (sendMessageMode == "async") {
      producer.setRetryTimesWhenSendAsyncFailed(retryTimesWhenSendAsyncFailed)
    }
  }

  override def process(): Status = {
    val channel = getChannel
    val tx = channel.getTransaction
    try {
      tx.begin()
      val event: Event = channel.take
      if (!(event == null || event.getBody == null || (event.getBody.length == 0))) {
        val msg = new Message(topic, tags, keys, flag, event.getBody, waitStoreMsgOK)
        sendMessageMode match {
          case "sync" =>
            // Send message in synchronous mode. This method returns only when the sending procedure totally completes.
            producer.send(msg)
          case "async" =>
            // Send Messages Asynchronously
            producer.send(msg, new SendCallback() {
              override def onException(e: Throwable): Unit = ???

              override def onSuccess(sendResult: SendResult): Unit = ???
            })
          case "oneway" =>
            // Send Messages in One-way Mode
            producer.sendOneway(msg)
          case _ =>
            throw new Exception("The value of the parameter sendMessageMode is set incorrectly.")
        }
      }
      tx.commit()
      Status.READY
    } catch {
      case _: Exception =>
        try {
          tx.rollback()
        } catch {
          case _: Exception =>
        }
        Status.BACKOFF
    } finally {
      tx.close()
    }
  }

  override def stop() = {
    producer.shutdown()
    super.stop()
  }
}
